#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
统计绘图工具
功能：读取数据文件，进行显著性差异分析，生成箱线图和P值热图，支持设置箱线图Y轴范围

版本: 1.3.0
作者: 陈振玺
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import argparse
import sys
from pathlib import Path
from scipy import stats
from scipy.stats import mannwhitneyu, ttest_ind, wilcoxon, ttest_rel
from itertools import combinations
import warnings
from matplotlib.patches import Rectangle

# 设置字体支持
plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Arial', 'sans-serif']
plt.rcParams['axes.unicode_minus'] = False


class StatisticalPlotter:
    """
    统计绘图工具类
    
    功能：
    - 读取CSV/TSV/XLSX数据文件
    - 进行两组数据的显著性差异分析
    - 生成箱线图和P值热图的组合图表
    - 支持多种统计检验方法
    - 支持数据排序和匹配验证
    
    版本: 1.3.0
    作者: 陈振玺
    """
    
    def __init__(self):
        """
        初始化统计绘图工具
        """
        self.data = None
        self.x_column = None
        self.y_column = None
        self.color_column = None  # 颜色列
        self.method = 'mannwhitney'  # 默认使用Mann-Whitney U检验
        self.reference_column = None
        self.y_min = None  # Y轴最小值
        self.y_max = None  # Y轴最大值
        self.available_methods = {
            'mannwhitney': '曼-惠特尼U检验（非参数）',
            'ttest': 't检验（参数）',
            'wilcoxon': '威尔科克森符号秩检验（配对非参数）',
            'paired_ttest': '配对t检验（配对参数）'
        }
        
    def load_data(self, file_path):
        """
        加载数据文件
        
        Args:
            file_path (str): 数据文件路径
            
        Returns:
            bool: 加载成功返回True，否则返回False
        """
        try:
            file_path = Path(file_path)
            
            if not file_path.exists():
                print(f"错误: 文件 '{file_path}' 不存在")
                return False
            
            # 根据文件扩展名选择读取方法
            if file_path.suffix.lower() == '.xlsx':
                self.data = pd.read_excel(file_path)
            elif file_path.suffix.lower() in ['.csv']:
                self.data = pd.read_csv(file_path)
            elif file_path.suffix.lower() in ['.tsv', '.txt']:
                self.data = pd.read_csv(file_path, sep='\t')
            else:
                print(f"错误: 不支持的文件格式 '{file_path.suffix}'")
                return False
            
            print(f"成功加载数据: {self.data.shape[0]} 行, {self.data.shape[1]} 列")
            print(f"列名: {list(self.data.columns)}")
            return True
            
        except Exception as e:
            print(f"加载数据失败: {e}")
            return False
    
    def validate_columns(self, x_column, y_column, color_column=None):
        """
        验证指定的列是否存在
        
        Args:
            x_column (str): X轴列名
            y_column (str): Y轴列名
            color_column (str): 颜色列名
            
        Returns:
            bool: 验证通过返回True，否则返回False
        """
        if self.data is None:
            print("错误: 请先加载数据")
            return False
        
        missing_columns = []
        if x_column not in self.data.columns:
            missing_columns.append(x_column)
        if y_column not in self.data.columns:
            missing_columns.append(y_column)
        if color_column and color_column not in self.data.columns:
            missing_columns.append(color_column)
        
        if missing_columns:
            print(f"错误: 以下列不存在: {missing_columns}")
            print(f"可用列: {list(self.data.columns)}")
            return False
        
        self.x_column = x_column
        self.y_column = y_column
        self.color_column = color_column
        return True
    
    def sort_data_by_reference(self, reference_column):
        """
        根据参考列对数据进行排序
        
        Args:
            reference_column (str): 参考列名
            
        Returns:
            bool: 排序成功返回True，否则返回False
        """
        if reference_column not in self.data.columns:
            print(f"错误: 参考列 '{reference_column}' 不存在")
            return False
        
        self.reference_column = reference_column
        self.data = self.data.sort_values(by=reference_column)
        print(f"数据已按 '{reference_column}' 列排序")
        return True
    
    def validate_paired_data(self):
        """
        验证配对数据的完整性
        
        Returns:
            bool: 验证通过返回True，否则返回False
        """
        if self.reference_column is None:
            return True  # 非配对分析，无需验证
        
        # 检查每个分组是否有相同的参考值
        groups = self.data.groupby(self.x_column)
        reference_sets = {}
        
        for group_name, group_data in groups:
            reference_sets[group_name] = set(group_data[self.reference_column].values)
        
        # 检查所有组的参考值是否完全匹配
        group_names = list(reference_sets.keys())
        if len(group_names) < 2:
            print("警告: 数据中少于2个分组")
            return False
        
        base_set = reference_sets[group_names[0]]
        for i in range(1, len(group_names)):
            if reference_sets[group_names[i]] != base_set:
                print(f"警告: 分组 '{group_names[0]}' 和 '{group_names[i]}' 的参考值不完全匹配")
                print(f"分组 '{group_names[0]}' 参考值数量: {len(base_set)}")
                print(f"分组 '{group_names[i]}' 参考值数量: {len(reference_sets[group_names[i]])}")
                return False
        
        print(f"配对数据验证通过，每组包含 {len(base_set)} 个配对样本")
        return True
    
    def perform_statistical_test(self, group1_data, group2_data, method='mannwhitney'):
        """
        执行统计检验
        
        Args:
            group1_data (array-like): 第一组数据
            group2_data (array-like): 第二组数据
            method (str): 统计检验方法
            
        Returns:
            float: P值
        """
        try:
            # 移除NaN值
            group1_clean = np.array(group1_data).astype(float)
            group2_clean = np.array(group2_data).astype(float)
            group1_clean = group1_clean[~np.isnan(group1_clean)]
            group2_clean = group2_clean[~np.isnan(group2_clean)]
            
            if len(group1_clean) == 0 or len(group2_clean) == 0:
                return np.nan
            
            if method == 'mannwhitney':
                # Mann-Whitney U检验（独立样本非参数检验）
                statistic, p_value = mannwhitneyu(group1_clean, group2_clean, alternative='two-sided')
            elif method == 'ttest':
                # 独立样本t检验
                statistic, p_value = ttest_ind(group1_clean, group2_clean)
            elif method == 'wilcoxon':
                # Wilcoxon符号秩检验（配对非参数检验）
                if len(group1_clean) != len(group2_clean):
                    print("警告: Wilcoxon检验要求配对数据长度相等")
                    return np.nan
                statistic, p_value = wilcoxon(group1_clean, group2_clean)
            elif method == 'paired_ttest':
                # 配对t检验
                if len(group1_clean) != len(group2_clean):
                    print("警告: 配对t检验要求配对数据长度相等")
                    return np.nan
                statistic, p_value = ttest_rel(group1_clean, group2_clean)
            else:
                print(f"错误: 不支持的统计方法 '{method}'")
                return np.nan
            
            return p_value
            
        except Exception as e:
            print(f"统计检验失败: {e}")
            return np.nan
    
    def calculate_pairwise_pvalues(self):
        """
        计算所有分组间的两两P值
        
        Returns:
            pandas.DataFrame: P值矩阵
        """
        if self.data is None or self.x_column is None or self.y_column is None:
            print("错误: 请先加载数据并指定X、Y列")
            return None
        
        # 获取所有分组并排序以确保一致性
        groups = self.data.groupby(self.x_column)
        group_names = sorted(groups.groups.keys())
        
        if len(group_names) < 2:
            print("错误: 至少需要2个分组进行比较")
            return None
        
        # 初始化P值矩阵
        p_matrix = pd.DataFrame(index=group_names, columns=group_names, dtype=float)
        
        # 计算两两比较的P值
        for i, group1 in enumerate(group_names):
            for j, group2 in enumerate(group_names):
                if i == j:
                    p_matrix.loc[group1, group2] = 1.0  # 自己与自己比较P值为1
                elif i < j:  # 只计算上三角矩阵，避免重复计算
                    group1_data = groups.get_group(group1)[self.y_column]
                    group2_data = groups.get_group(group2)[self.y_column]
                    
                    if self.reference_column and self.method in ['wilcoxon', 'paired_ttest']:
                        # 配对检验：需要按参考列匹配数据
                        group1_df = groups.get_group(group1).set_index(self.reference_column)
                        group2_df = groups.get_group(group2).set_index(self.reference_column)
                        
                        # 找到共同的参考值
                        common_refs = group1_df.index.intersection(group2_df.index)
                        if len(common_refs) == 0:
                            p_value = np.nan
                        else:
                            group1_matched = group1_df.loc[common_refs, self.y_column]
                            group2_matched = group2_df.loc[common_refs, self.y_column]
                            p_value = self.perform_statistical_test(group1_matched, group2_matched, self.method)
                    else:
                        # 独立样本检验
                        p_value = self.perform_statistical_test(group1_data, group2_data, self.method)
                    
                    p_matrix.loc[group1, group2] = p_value
                    p_matrix.loc[group2, group1] = p_value  # 对称矩阵
        
        return p_matrix
    
    def create_boxplot(self, ax):
        """
        创建箱线图
        
        Args:
            ax (matplotlib.axes.Axes): 绘图轴对象
        """
        # 使用seaborn绘制箱线图
        if self.color_column:
            # 根据颜色列设置箱线图颜色
            sns.boxplot(data=self.data, x=self.x_column, y=self.y_column, hue=self.color_column, ax=ax, notch=True)
        else:
            # 使用默认颜色
            sns.boxplot(data=self.data, x=self.x_column, y=self.y_column, ax=ax, notch=True)
        
        # 设置图表属性
        ax.set_xlabel(f'{self.x_column}', fontsize=12, fontweight='bold')
        ax.set_ylabel(f'{self.y_column}', fontsize=12, fontweight='bold')
        ax.set_title('Box Plot', fontsize=14, fontweight='bold')
        ax.grid(True, alpha=0.3)
        
        # 旋转x轴标签90度以提高可读性，并调整字体大小
        ax.tick_params(axis='x', rotation=90, labelsize=10)
        # 调整底部边距以防止标签被截断
        ax.margins(x=0.01)
        
        # 设置Y轴范围
        if self.y_min is not None or self.y_max is not None:
            y_min = self.y_min if self.y_min is not None else ax.get_ylim()[0]
            y_max = self.y_max if self.y_max is not None else ax.get_ylim()[1]
            ax.set_ylim(y_min, y_max)
    
    def create_pvalue_heatmap(self, ax, p_matrix):
        """
        创建P值热图
        
        Args:
            ax (matplotlib.axes.Axes): 绘图轴对象
            p_matrix (pandas.DataFrame): P值矩阵
        """
        # 创建热图
        im = ax.imshow(p_matrix.values, cmap='RdYlBu_r', aspect='auto', vmin=0, vmax=1)
        
        # 设置坐标轴标签
        ax.set_xticks(range(len(p_matrix.columns)))
        ax.set_yticks(range(len(p_matrix.index)))
        # 将X轴标签旋转90度垂直显示，调整字体大小和对齐方式
        ax.set_xticklabels(p_matrix.columns, rotation=90, ha='center', fontsize=10)
        ax.set_yticklabels(p_matrix.index, fontsize=10)
        
        # 在每个格子中添加P值文本
        for i in range(len(p_matrix.index)):
            for j in range(len(p_matrix.columns)):
                p_val = p_matrix.iloc[i, j]
                if pd.isna(p_val):
                    text = 'N/A'
                    color = 'black'
                elif p_val < 0.001:
                    text = '***'
                    color = 'white' if p_val < 0.5 else 'black'
                elif p_val < 0.01:
                    text = '**'
                    color = 'white' if p_val < 0.5 else 'black'
                elif p_val < 0.05:
                    text = '*'
                    color = 'white' if p_val < 0.5 else 'black'
                else:
                    text = f'{p_val:.3f}'
                    color = 'white' if p_val < 0.5 else 'black'
                
                ax.text(j, i, text, ha='center', va='center', color=color, fontweight='bold', fontsize=8)
        
        # 设置标题
        ax.set_title('Statistical Significance P-value Matrix', fontsize=14, fontweight='bold')
        
        # 添加颜色条
        cbar = plt.colorbar(im, ax=ax, shrink=0.8)
        cbar.set_label('P-value', rotation=270, labelpad=20, fontsize=12)
    
    def generate_plot(self, output_path, figsize=(16, 8)):
        """
        生成组合图表
        
        Args:
            output_path (str): 输出文件路径
            figsize (tuple): 图表大小
            
        Returns:
            str: 输出文件路径
        """
        if self.data is None:
            raise ValueError("请先加载数据")
        
        if self.x_column is None or self.y_column is None:
            raise ValueError("请先指定X、Y列")
        
        # 计算P值矩阵
        print(f"使用 {self.available_methods[self.method]} 进行显著性分析...")
        p_matrix = self.calculate_pairwise_pvalues()
        
        if p_matrix is None:
            raise ValueError("P值计算失败")
        
        # 创建图表
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
        
        # 绘制箱线图
        self.create_boxplot(ax1)
        
        # 绘制P值热图
        self.create_pvalue_heatmap(ax2, p_matrix)
        
        # 添加总标题，预留足够的顶部空间避免与子图标题重叠
        method_names_en = {
            'mannwhitney': 'Mann-Whitney U Test',
            'ttest': 'Independent t-test',
            'wilcoxon': 'Wilcoxon Signed-rank Test',
            'paired_ttest': 'Paired t-test'
        }
        fig.suptitle(f'Statistical Analysis Results - {method_names_en[self.method]}', 
                    fontsize=16, fontweight='bold', y=0.95)
        
        # 调整布局，增加子图间距和边距以提升可读性
        plt.tight_layout()
        # 设置更大的子图间距：wspace控制水平间距，hspace控制垂直间距
        # top参数为主标题预留空间，bottom和right为轴标签预留空间
        # 将wspace从0.4增加到0.8，彻底避免子图重叠问题
        plt.subplots_adjust(top=0.88, bottom=0.15, right=0.85, wspace=0.8, hspace=0.3)
        
        # 保存图表
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        
        plt.savefig(output_path, dpi=300, bbox_inches='tight', facecolor='white')
        plt.close()
        
        print(f"图表已保存到: {output_path}")
        return str(output_path)
    
    def print_summary(self):
        """
        打印数据摘要信息
        """
        if self.data is None:
            print("无数据可显示")
            return
        
        print("\n=== 数据摘要 ===")
        print(f"数据形状: {self.data.shape}")
        print(f"X列 ({self.x_column}) 分组:")
        
        groups = self.data.groupby(self.x_column)
        for name, group in groups:
            y_data = group[self.y_column].dropna()
            print(f"  {name}: {len(y_data)} 个样本, 均值={y_data.mean():.3f}, 标准差={y_data.std():.3f}")
        
        if self.reference_column:
            print(f"\n参考列: {self.reference_column}")
            print(f"统计方法: {self.available_methods[self.method]}")
        else:
            print(f"\n统计方法: {self.available_methods[self.method]}")


def main():
    """
    命令行主函数
    """
    parser = argparse.ArgumentParser(
        description='统计绘图工具 - 生成箱线图和显著性差异分析图',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  %(prog)s -s data.csv -x group -y value -o result.png
  %(prog)s -s data.xlsx -x treatment -y score -m ttest -o analysis.pdf
  %(prog)s -s data.tsv -x condition -y measurement -m wilcoxon -r sample_id -o paired_analysis.png
  
支持的统计方法:
  mannwhitney  : Mann-Whitney U检验（非参数，独立样本）
  ttest        : t检验（参数，独立样本）
  wilcoxon     : Wilcoxon符号秩检验（非参数，配对样本）
  paired_ttest : 配对t检验（参数，配对样本）
        """
    )
    
    parser.add_argument(
        '-s', '--source',
        required=True,
        help='输入数据文件路径（支持CSV、TSV、XLSX格式）'
    )
    
    parser.add_argument(
        '-o', '--output',
        required=True,
        help='输出图表文件路径（支持PNG、PDF、SVG等格式）'
    )
    
    parser.add_argument(
        '-x', '--x-column',
        required=True,
        help='X轴列名（分组列）'
    )
    
    parser.add_argument(
        '-y', '--y-column',
        required=True,
        help='Y轴列名（数值列）'
    )
    
    parser.add_argument(
        '-m', '--method',
        choices=['mannwhitney', 'ttest', 'wilcoxon', 'paired_ttest'],
        default='mannwhitney',
        help='统计检验方法（默认: mannwhitney）'
    )
    
    parser.add_argument(
        '-r', '--reference',
        help='参考列名（用于配对分析时的样本匹配）'
    )
    
    parser.add_argument(
        '-c', '--colormap',
        default='Set3',
        help='颜色映射方案（默认: Set3）'
    )
    
    parser.add_argument(
        '--color-column',
        help='颜色列名（用于根据指定列的值设置箱线图颜色）'
    )
    
    parser.add_argument(
        '--figsize',
        nargs=2,
        type=float,
        default=[16, 8],
        help='图表大小 宽度 高度（默认: 16 8）'
    )
    
    parser.add_argument(
        '--y-min',
        type=float,
        help='箱线图Y轴最小值'
    )
    
    parser.add_argument(
        '--y-max',
        type=float,
        help='箱线图Y轴最大值'
    )
    
    parser.add_argument(
        '--version',
        action='version',
        version='统计绘图工具 1.0.0 (作者: 陈振玺)'
    )
    
    args = parser.parse_args()
    
    try:
        # 创建绘图工具实例
        plotter = StatisticalPlotter()
        
        # 加载数据
        if not plotter.load_data(args.source):
            return 1
        
        # 验证列
        if not plotter.validate_columns(args.x_column, args.y_column, args.color_column):
            return 1
        
        # 设置统计方法
        plotter.method = args.method
        
        # 设置Y轴范围
        plotter.y_min = args.y_min
        plotter.y_max = args.y_max
        
        # 处理参考列
        if args.reference:
            if not plotter.sort_data_by_reference(args.reference):
                return 1
            
            # 验证配对数据
            if args.method in ['wilcoxon', 'paired_ttest']:
                if not plotter.validate_paired_data():
                    print("警告: 配对数据验证失败，但将继续分析")
        
        # 打印数据摘要
        plotter.print_summary()
        
        # 生成图表
        output_file = plotter.generate_plot(args.output, tuple(args.figsize))
        
        print(f"\n分析完成!")
        print(f"输出文件: {output_file}")
        
    except Exception as e:
        print(f"错误: {e}")
        return 1
    
    return 0


if __name__ == '__main__':
    exit(main())